#pragma once
#include "LastBlockEncryption.h"
#include "LastBlockDecryption.h"

namespace AppliedCryptography
{
	namespace AC = AppliedCryptography;

	template<unsigned int TKeySize, unsigned int TBlockSize>
	class ChypertextStealing
	{
	public:
		class Encryption : public AC::LastBlockEncryption<TKeySize, TBlockSize>
		{
			public:
				virtual AC::Block<TBlockSize> Process(
					const AC::Cryptor<TKeySize,TBlockSize>& cryptor, 
					const AC::Block<TBlockSize>& last /* Plaintext */, 
					unsigned int lastBlockLength,
					AC::Block<TBlockSize>& secondToLast /* Should be encrypted once */) const
				{
					AC::Block<TBlockSize> rblock;
					AC::Block<TBlockSize> xord;

					for(unsigned int i=0; i < lastBlockLength; i++)
						rblock[i] = secondToLast[i];

					xord = secondToLast ^ last;

					cryptor.Process(xord, secondToLast);

					return rblock;
				};
		};

		class Decryption : public AC::LastBlockDecryption<TKeySize, TBlockSize>
		{
			public:
				virtual AC::Block<TBlockSize> Process(
					const AC::Cryptor<TKeySize,TBlockSize>& cryptor, 
					const AC::Block<TBlockSize>& last /* Plaintext */, 
					unsigned int lastBlockLength,
					AC::Block<TBlockSize>& secondToLast /* Should be encrypted once */) const
				{
					const AC::Cryptor<TKeySize,TBlockSize>& dec = cryptor;
					AC::Block<TBlockSize> rblock;
					AC::Block<TBlockSize> xord;
					AC::Block<TBlockSize> secondCrypted;

					// Step #3
					xord = secondToLast ^ last;

					// Step #4
					for(unsigned int i=0; i < lastBlockLength; i++)
						rblock[i] = xord[i];
					
					// Step #5
					secondCrypted = last;
					for(unsigned int i=0; i < TBlockSize-lastBlockLength; i++)
						secondCrypted[lastBlockLength+i] = xord[lastBlockLength+i];

					// Step #6
					dec.Process(secondCrypted, secondToLast);

					return rblock;
				};
		};
	};
};
